import os
import pandas as pd
import numpy as np
from pathlib import Path
from shutil import rmtree, move
import torch

import cv2
import toml
import numpy as np
from PIL import Image
from albumentations.pytorch import ToTensorV2
from albumentations import(
    OneOf, Resize, Normalize, Compose, Transpose,
    HorizontalFlip, VerticalFlip, Flip, Cutout, RandomCrop,
    CenterCrop, ShiftScaleRotate, Rotate, RandomSizedCrop,
    RandomContrast, RandomBrightness, RandomBrightnessContrast,
    RandomGamma, CLAHE, IAASharpen, IAAEmboss, FancyPCA, Sharpen, IAAPerspective,
    GaussianBlur, GaussNoise, Blur, MotionBlur, MedianBlur, 
    HueSaturationValue, OpticalDistortion, GridDistortion, RandomRotate90, ChannelShuffle, RGBShift
)
import torchvision.transforms


def get_transforms(data_aug_config, img_size, mean=[1,1,1], std=[1,1,1], padding=16):
    transform_list = []
    if len(data_aug_config) == 0:
        return transform_list
    
    x = sorted(data_aug_config.items(), key=lambda x: x[1]['order'])
    for name, param in x:
        if not param['use']:
            continue
        if name == 'Resize':
            transform_list.append(Resize(height=img_size[0], width=img_size[1]))
        elif name == 'RandomRotate90': 
            transform_list.append(RandomRotate90(p=param['p']))
        elif name == 'VerticalFlip': 
            transform_list.append(VerticalFlip(p=param['p']))
        elif name == 'HorizontalFlip': 
            transform_list.append(HorizontalFlip(p=param['p']))
        elif name == 'Transpose':
            transform_list.append(Transpose(p=param['p']))
        elif name == 'IAAPerspective': 
            transform_list.append(IAAPerspective(p=param['p']))
        elif name == 'Rotate': 
            transform_list.append(Rotate(limit=param['limit'], p=param['p']))
        elif name == 'RandomSizedCrop': 
            transform_list.append(RandomSizedCrop(p=param['p'],min_max_height=(int(img_size[0] * 0.9), img_size[0]), height=img_size[0], width=img_size[1]))
        elif name == 'RandomBrightnessContrast': 
            transform_list.append(RandomBrightnessContrast(p=param['p'], brightness_limit=param['brightness_limit'], contrast_limit=param['contrast_limit']))
        elif name == 'ChannelShuffle': 
            transform_list.append(ChannelShuffle(p=param['p']))
        elif name == 'RGBShift': 
            transform_list.append(RGBShift(p=param['p'], r_shift_limit=param['r_shift_limit'], g_shift_limit=param['g_shift_limit'], b_shift_limit=param['b_shift_limit']))
        elif name == 'CLAHE': 
            transform_list.append(CLAHE(p=param['p'], clip_limit=param['clip_limit'], tile_grid_size=param['tile_grid_size']))

            # CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), always_apply=False, p=0.5)
        elif name == 'Normalize': 
            transform_list.append(Normalize(p=1.0, mean=mean, std=std, max_pixel_value=255.0))
        else:
            pass
    
    return transform_list

def multi_transforms_medium_4(transforms_list):

    transforms_list.append(ToTensorV2())
    transforms = Compose(transforms_list)
    return transforms

data_aug_func = {"Resize":{"order":1, "use":False}, 
                     "RandomRotate90":{"order":2, "use":False, "p":0.5},
                     "VerticalFlip":{"order":3, "use":False, "p":0.5}, 
                     "HorizontalFlip":{"order":4, "use":False, "p":0.5},
                     "Transpose":{"order":5, "use":False, "p":0.5}, 
                     "IAAPerspective":{"order":6, "use":False, "p":0.3},
                     "Rotate":{"order":7, "use":False, "p":0.5, "limit":5},
                     "RandomSizedCrop":{"order":8, "use":False, "p":0.5},
                     "RandomBrightnessContrast":{"order":13, "use":True, "p":0.9, "brightness_limit":0.0, "contrast_limit":0.8},
                     "ChannelShuffle":{"order":11, "use":False, "p":0.7},
                     "RGBShift":{"order":12, "use":False, "p":0.5, "r_shift_limit":50, "g_shift_limit":50, "b_shift_limit":50},
                     "CLAHE":{"order":9, "use":True, "p":0.5, "clip_limit":4, "tile_grid_size":(8, 8)},
                     "Normalize":{"order":14, "use":False},
                    }

transform_func = get_transforms(data_aug_config=data_aug_func, img_size=(1000, 1000))
train_transforms = multi_transforms_medium_4(transform_func)


def cal_img(func, img, path, dst_path, times=10):
    imgPath = os.path.join(path, img)
    image = cv2.imread(imgPath)
    # print(image.shape)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    for i in range(times):
        img_out = func(image=image)['image'].numpy()
        img_out = img_out.transpose(1,2,0)
        img_out = cv2.cvtColor(img_out, cv2.COLOR_BGR2RGB)
    
        img_new_name = img[:-4] + '_%d'%i + '.jpg'
    # print(img_out.shape)
        cv2.imwrite(os.path.join(dst_path, img_new_name), img_out)

cal_img(train_transforms, img=r'18850_HCA434501AB_TEAOH520_70_-425.009_-96.699__4.32839_20230420_190354-Normal_.jpg', 
                        path=r'/opt/data/private/project/T5/12.9/18853_other/data/TNSDO1', 
                        dst_path=r'/opt/data/private/project/T5/12.9/18853_other/data/TNSDO1', 
                        times=50)

# cal_img(train_transforms, img=r'18850_HCA435D01CB_TEAOH520_74_-661.104_888.068_M_10.0009_20230522_083156.jpg', path=r'F:\T5\12.9\data\18853new2\data\Images_Other\TNNNF0', dst_path=r'F:\T5\12.9\data\18853new2\data\Images_Other\TNNNF0_aug', times=20)